#include <time.h>
#include <cmath>
#include <fstream>
#include <iostream>
#include <random>

// glog
#include "glog/logging.h"

// gtsam
#include "gtsam/base/Matrix.h"
#include "gtsam/base/Vector.h"
#include "gtsam/inference/Key.h"
#include "gtsam/inference/Symbol.h"
#include "gtsam/nonlinear/ISAM2.h"
#include "gtsam/nonlinear/NonlinearFactorGraph.h"
#include "gtsam/nonlinear/Values.h"

// unknown variables curve params: gtsam::Vector3
// measurement error: y_pred - y_measurement
class CurveFitting : public gtsam::NoiseModelFactor1<gtsam::Vector3> {
 public:
  CurveFitting(gtsam::Key param_key,
               const double x,
               const double y,
               const gtsam::SharedNoiseModel& model)
      : gtsam::NoiseModelFactor1<gtsam::Vector3>(model, param_key),
        x_(x),
        y_(y) {}

  gtsam::Vector evaluateError(
      const gtsam::Vector3& params,
      boost::optional<gtsam::Matrix&> H = boost::none) const {
    if (H) {
      (*H) = (gtsam::Matrix(1, 3) << x_ * x_, x_, 1.0).finished();
    }
    double error = params[0] * x_ * x_ + params[1] * x_ + params[2] - y_;
    return (gtsam::Vector(1) << error).finished();
  }

 private:
  double x_;
  double y_;
};

int main(int argc, char** argv) {
  // random gen
  uint32_t curr_time = static_cast<uint32_t>(std::time(nullptr));
  std::random_device rd{};
  std::mt19937 gen{rd()};
  LOG(INFO) << "random generator seed: " << curr_time;
  gen.seed(curr_time);
  std::normal_distribution<double> normal_dist{0.0, 5};
  // generate ground truth dataset
  std::vector<double> x_datas, y_datas;
  double params[] = {0.0031221323, 0.21532324, 4.3546546546};

  std::ofstream data_fout("/tmp/sim_data.txt");
  data_fout.precision(10);
  for (int i = -100; i < 200; i++) {
    x_datas.push_back(i);
    double y_noise =
        params[0] * i * i + params[1] * i + params[2] + normal_dist(gen);
    y_datas.push_back(y_noise);
    data_fout << i << " " << y_noise << std::endl;
  }
  data_fout.close();

  // isam
  gtsam::ISAM2Params isam2_params;
  isam2_params.relinearizeSkip = 100;
  isam2_params.relinearizeThreshold = 0.1;
  gtsam::ISAM2 isam2_solver(isam2_params);

  gtsam::NonlinearFactorGraph graph;
  gtsam::Values init_estimates;
  gtsam::Values results;

  // data noise model
  auto data_noise_model =
      gtsam::noiseModel::Diagonal::Sigmas(gtsam::Vector1(5));
  // build factor graph
  for (int i = 0; i < y_datas.size(); i++) {
    graph.add(boost::make_shared<CurveFitting>(
        gtsam::Symbol('x', 0), x_datas[i], y_datas[i], data_noise_model));
  }
  // init vaules for params
  init_estimates.insert<gtsam::Vector3>(gtsam::Symbol('x', 0),
                                        gtsam::Vector3(0.0, 0.0, 0.0));
  // optimization
  isam2_solver.update(graph, init_estimates);
  results = isam2_solver.calculateEstimate();

  LOG(INFO) << "results: "
            << results.at<gtsam::Vector3>(gtsam::Symbol('x', 0)).transpose();
}